#include "RBTree.h"

RBTreeNode *NIL;
rbtree_compare g_compare;

void rbtree_insert_fixup(RBTreeNode **root,RBTreeNode *node);
void rbtree_right_rotate(RBTreeNode **root,RBTreeNode *node);
void rbtree_left_rotate(RBTreeNode **root,RBTreeNode *node);
void rbtree_transplant(RBTreeNode **root,RBTreeNode *parent,RBTreeNode *child);
RBTreeNode *rbtree_rear(RBTreeNode *node);
void rbtree_delete_fixup(RBTreeNode **root,RBTreeNode *node);
RBTreeNode *rbtree_create_node(void *key,void *value);

RBTreeNode *rbtree_create_node(void *key, void *value)
{
    RBTreeNode *node = (RBTreeNode *)malloc(sizeof(RBTreeNode));
    if (node == NULL)
    {
        return NULL;
    }
    node->color = RBT_RED;
    node->pkey = *((void**)key); 
    node->pvalue = value;
    node->pLChild = NIL; 
    node->pRChild = NIL;
    node->pParentNode = NIL;
    return node;
}

int rbtree_insert_node(RBTreeNode **root,void *key, void *value)
{
    if (root == NULL || key == NULL || value == NULL || g_compare == NULL)
    {
        return -1;
    }
    
    RBTreeNode * node = rbtree_create_node(key, value);
    if (node == NULL)
    {
        return -1;
    }

    //找插入位置
    RBTreeNode *parent = NIL;
    RBTreeNode *pos = *root;
    while (pos != NIL)
    {
        parent = pos;
        if (g_compare(key,&(pos->pkey)) < 0)
        {
            pos = pos->pLChild;
        }
        else if (g_compare(key,&(pos->pkey)) > 0)
        {
            pos = pos->pRChild;
        }
        else if (g_compare(key,&(pos->pkey)) == 0)
        {
            return 1;
        }
    }

    node->pParentNode = parent;
    if (parent == NIL)
    {
        *root = node;
    }
    else if (g_compare(key,&parent->pkey) < 0)
    {
        parent->pLChild = node;
    }
    else
    {
        parent->pRChild = node;
    }
    
    rbtree_insert_fixup(root,node);

    return 0;
}

void rbtree_insert_fixup(RBTreeNode **root, RBTreeNode *node)
{
    while (node->pParentNode->color == RBT_RED)
    {
        if (node->pParentNode->pParentNode->pLChild == node->pParentNode)//父亲是左孩子
        {
            RBTreeNode *uncle = node->pParentNode->pParentNode->pRChild;
            if (uncle->color == RBT_RED)
            {
                node->pParentNode->color = RBT_BLACK;
                uncle->color = RBT_BLACK;
                node->pParentNode->pParentNode->color = RBT_RED;
                node = node->pParentNode->pParentNode;
            }
            else
            {
                if (node->pParentNode->pRChild == node)
                {
                    node = node->pParentNode;
                    rbtree_left_rotate(root, node);
                }
                node->pParentNode->color = RBT_BLACK;
                node->pParentNode->pParentNode->color = RBT_RED;
                // rbtree_right_rotate(root,uncle->pParentNode); //花好久找这个坑，阿啊啊啊啊
                rbtree_right_rotate(root,node->pParentNode->pParentNode);
            }
        }
        else//父亲是右孩子
        {
            RBTreeNode *uncle = node->pParentNode->pParentNode->pLChild;
            if (uncle->color == RBT_RED)
            {
                node->pParentNode->color = RBT_BLACK;
                uncle->color = RBT_BLACK;
                node->pParentNode->pParentNode->color = RBT_RED;
                node = node->pParentNode->pParentNode;
            }
            else
            {
                if (node->pParentNode->pLChild == node)
                {
                    node = node->pParentNode;
                    rbtree_right_rotate(root, node);
                }
                node->pParentNode->color = RBT_BLACK;
                // uncle->pParentNode->color = RBT_RED;
                node->pParentNode->pParentNode->color = RBT_RED;//这里不能按上面写，因为对于NIL结点，它是找不到父亲的
                rbtree_left_rotate(root,node->pParentNode->pParentNode);
            }
        }
    }
    (*root)->color = RBT_BLACK;
}

void rbtree_right_rotate(RBTreeNode **root,RBTreeNode *node)
{
    RBTreeNode *lchid = node->pLChild;
    node->pLChild = lchid->pRChild;
    if (lchid->pRChild != NIL)
    {
        lchid->pRChild->pParentNode = node;
    }

    lchid->pParentNode = node->pParentNode;
    if (node->pParentNode == NIL)
    {
        *root = lchid;
    }
    else if (node->pParentNode->pLChild == node)
    {
        lchid->pParentNode->pLChild = lchid;
    }
    else
    {
        lchid->pParentNode->pRChild = lchid;
    }
    lchid->pRChild = node;
    node->pParentNode = lchid;
}

void rbtree_left_rotate(RBTreeNode **root, RBTreeNode *node)
{
    RBTreeNode* rchild = node->pRChild;
    node->pRChild = rchild->pLChild;
    if (rchild->pLChild != NIL)
    {
        rchild->pLChild->pParentNode = node;
    }

    rchild->pParentNode = node->pParentNode;
    if (node->pParentNode == NIL)
    {
        *root = rchild;
    }
    else if (node->pParentNode->pLChild == node)
    {
        node->pParentNode->pLChild = rchild;
    }
    else
    {
        node->pParentNode->pRChild = rchild;
    }
    rchild->pLChild = node;
    node->pParentNode = rchild;
    
}

int rbtree_delete_node(RBTreeNode **root, RBTreeNode *node)
{
    if (root == NULL || node == NULL || node == NIL)
    {
        return -1;
    }
    
    RBTreeNode *dnode = node;
    RBTreeNode *replaceNode;
    int color = node->color;

    if (node->pLChild == NIL)
    {
        replaceNode = node->pRChild;
        rbtree_transplant(root,node,node->pRChild);
    }
    else if (node->pRChild == NIL)
    {
        replaceNode = node->pLChild;
        rbtree_transplant(root,node,node->pLChild);
    }
    else
    {
        dnode = rbtree_rear(node);
        replaceNode = dnode->pRChild;
        color = dnode->color;

        if (dnode == node->pRChild)
        {
            replaceNode->pParentNode = dnode;
        }
        else
        {
            rbtree_transplant(root,dnode,replaceNode);
            dnode->pRChild = node->pRChild;
            dnode->pRChild->pParentNode = dnode;
        }

        rbtree_transplant(root,node,dnode);
        dnode->pLChild = node->pLChild;
        node->pLChild->pParentNode = dnode;
        dnode->color = node->color;
    }

    free(node);
    
    if (color == RBT_BLACK)
    {
        rbtree_delete_fixup(root,replaceNode);
    }
    return 0;
}

void rbtree_transplant(RBTreeNode **root,RBTreeNode *parent, RBTreeNode *child)
{
    if (parent->pParentNode == NIL)
    {
        *root = child;
    }
    else if (parent->pParentNode->pLChild == parent)
    {
        parent->pParentNode->pLChild = child;
    }
    else
    {
        parent->pParentNode->pRChild = child;
    }
    child->pParentNode = parent->pParentNode;
}

RBTreeNode *rbtree_rear(RBTreeNode *node)
{
    RBTreeNode* child = node->pRChild;
    while (child->pLChild != NIL)
    {
        child = child->pLChild;
    }
    return child;
}

void rbtree_delete_fixup(RBTreeNode **root, RBTreeNode *node)
{
    while (*root != node && node->color == RBT_BLACK)
    {
        if (node == node->pParentNode->pLChild)
        {
            RBTreeNode *sibling = node->pParentNode->pRChild;
            if (sibling->color == RBT_RED)//兄弟是红色
            {
                sibling->color = RBT_BLACK;
                node->pParentNode->color = RBT_RED;
                rbtree_left_rotate(root,node->pParentNode);
                sibling = node->pParentNode->pRChild;
            }
            else if (sibling->pLChild->color == RBT_BLACK && sibling->pRChild->color == RBT_BLACK)//兄弟及后代都是黑
            {
                sibling->color = RBT_RED;
                node = node->pParentNode;
            }
            else
            {
                if (sibling->pLChild->color == RBT_RED)
                {
                    sibling->color = RBT_RED;
                    sibling->pLChild->color = RBT_BLACK;
                    rbtree_right_rotate(root, sibling);
                    sibling = node->pParentNode->pRChild;
                }
                sibling->color = node->pParentNode->color;
                sibling->color = RBT_BLACK;
                sibling->pRChild->color = RBT_BLACK;
                rbtree_left_rotate(root, node->pParentNode);
                node = *root;
            }
        }
        else
        {
            RBTreeNode *sibling = node->pParentNode->pLChild;
            if (sibling->color == RBT_RED)//兄弟是红色
            {
                sibling->color = RBT_BLACK;
                node->pParentNode->color = RBT_RED;
                rbtree_right_rotate(root,node->pParentNode);
                sibling = node->pParentNode->pLChild;
            }
            else if (sibling->pLChild->color == RBT_BLACK && sibling->pRChild->color == RBT_BLACK)//兄弟及后代都是黑
            {
                sibling->color = RBT_RED;
                node = node->pParentNode;
            }
            else
            {
                if (sibling->pRChild->color == RBT_RED)
                {
                    sibling->color = RBT_RED;
                    sibling->pRChild->color = RBT_BLACK;
                    rbtree_left_rotate(root, sibling);
                    sibling = node->pParentNode->pLChild;
                }
                sibling->color = node->pParentNode->color;
                sibling->color = RBT_BLACK;
                sibling->pLChild->color = RBT_BLACK;
                rbtree_right_rotate(root, node->pParentNode);
                node = *root;
            }
        }
    }
    node->color = RBT_BLACK;
}

RBTreeNode *rbtree_search(RBTreeNode *root, void *key)
{
    if (g_compare == NULL)
    {
        return NIL;
    }
    
    RBTreeNode *x = root;

    while (x != NIL && g_compare(&x->pkey,key) != 0)
    {
        if (g_compare(&x->pkey,key) > 0)
        {
            x = x->pLChild;
        }
        else if(g_compare(&x->pkey,key) < 0)
        {
            x = x->pRChild;
        }
    }
    return x;
}

void rbtree_initNIL()
{
    NIL = (RBTreeNode *)malloc(sizeof(RBTreeNode));
    if (NIL == NULL)
    {
        return;
    }
    NIL->color = RBT_BLACK;
    NIL->pkey = NIL;
    NIL->pvalue = NIL;
    NIL->pLChild = NIL;
    NIL->pRChild = NIL;
    NIL->pParentNode = NIL;

    g_compare = NULL;
}
